import matplotlib.colors
import pandas as pd
import numpy as np
import os
import ast

import torch
from openai import OpenAI
import tiktoken
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA

csv = pd.read_csv("./MNIST/fine_food_reviews_1k.csv", index_col=0)
data = csv[["Time", "ProductId", "UserId", "Score", "Summary", "Text"]]
data = data.dropna()
data["Combined"] = "Summary:" + data.Summary.str.strip() +  ";Text:" + data.Text.str.strip()
print(data.head(3))
# 指定分词器
tokenizer = tiktoken.get_encoding("cl100k_base")



client = OpenAI(
    api_key="77140bfa516bad561310adc020fe67c9.txZAmho2OQBC7w1z",
    base_url="https://open.bigmodel.cn/api/paas/v4/"
)
def embedding_text(text):
    res = client.embeddings.create(input=text, model="embedding-2")
    return res.data[0].embedding

print(embedding_text("智谱调用测试"))
"""
data["embedding"] = data["Combined"].apply(embedding_text)
data.to_csv("./MNIST/out.csv")
"""

plt.rcParams["font.sans-serif"]=["SimHei"] #"宋体黑体" 支持中文
plt.rcParams["axes.unicode_minus"]=False
data = pd.read_csv("./MNIST/embedding_output_1k.csv")
print(f"data['embedding'][0]:{data['embedding'][0]}")
print(f"type(data['embedding'][0]):{type(data['embedding'][0])}")
data["embedding_vec"] = data['embedding'].apply(ast.literal_eval)
print(f"data['embedding_vec'][0]:{data['embedding_vec'][0]}")
print(f"type(data['embedding_vec'][0]):{type(data['embedding_vec'][0])}")

matrix = np.vstack(data['embedding_vec'].values)
print(f"matrix.shape:{matrix.shape}")
print(f"matrix:{matrix}")
# 降维度方法 设置维2维
"""
tsne = TSNE(n_components=2)
matrix_2d = tsne.fit_transform(matrix)
"""
matrix_2d = PCA(n_components=2).fit_transform(matrix)
print(f"matrix_2d.shape:{matrix_2d.shape}")

colormap = matplotlib.colors.ListedColormap(["red", "darkorange", "gold", "turquoise", "darkgreen"])
plt.scatter(x=matrix_2d[:, 0], y=matrix_2d[:, 1], c=data["Score"].values-1, cmap=colormap, alpha=0.3)

km = KMeans(3, init="k-means++")
km.fit(matrix)
labels = km.labels_
for i, label in enumerate(labels):
    plt.text(matrix_2d[i][0], matrix_2d[i][1], label)
plt.title("使用TSNE降维后的评论分布")
plt.show()